import pandas as pd

from app_config import get_engine_ts
from app_config import get_pro
import os

from datetime import datetime
from dateutil.relativedelta import relativedelta

""" 沪深300
"""
def offset_date_3M(original_date_str):
    # 将输入的日期字符串转换为日期对象
    original_date = datetime.strptime(original_date_str, '%Y%m%d')

    # 计算6个月后的日期
    new_date = original_date + relativedelta(months=3)

    # 将结果日期格式化为 'yyyyMMdd' 字符串
    return new_date.strftime('%Y%m%d')


def calc_hs300(strStartDate='', strEndDate='', str_filePre='', cur_000300Date='', pre_000300Date=''):
    if not os.path.exists(str_filePre):
        os.makedirs(str_filePre)
        print(f"文件夹 '{str_filePre}' 创建成功")
    engine = get_engine_ts()

    # 同中证全指
    df_stock_basic = get_pro().index_weight(index_code='000010.SH', start_date=cur_000300Date, end_date=cur_000300Date)
    df_stock_basic.to_excel(str_filePre + '上证180.xlsx')

    list_index_ts_code = df_stock_basic['con_code'].tolist()
    query_hfd = f"""
            SELECT ts_code, symbol, name, market, industry, list_date FROM `stock_basic`
            WHERE ts_code IN ({','.join(f"'{code}'" for code in list_index_ts_code)})
            AND market != '北交所'
            """
    df_stock_basic = pd.read_sql_query(query_hfd, engine)

    filter_ts_code = df_stock_basic['ts_code'].tolist()
    # 构建查询语句
    query = f"""
    SELECT ts_code,trade_date, amount FROM `daily`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
    AND trade_date >= '{strStartDate}'
    AND trade_date <= '{strEndDate}'
    """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 1000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(df_stock_basic, grouped_df, on='ts_code', how='left')

    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, total_mv FROM `daily_basic`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
        AND trade_date >= '{strStartDate}'
        AND trade_date <= '{strEndDate}'
        """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', '日均总市值']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')
    pd_merge.to_excel(str_filePre + '日均成交额.xlsx')

    # 按日均成交额 列降序排序
    df_sorted = pd_merge.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    df_sorted['i_日均成交'] = df_sorted.index.astype(int) + 1

    df_sorted = df_sorted.sort_values(by='i_日均成交').reset_index(drop=True)


    top_50_percent = df_sorted.sort_values('日均总市值', ascending=False).reset_index(drop=True)
    top_50_percent['i_日均市值'] = top_50_percent.index
    top_50_percent.to_excel(str_filePre + '沪深300_top50.xlsx')
    df_result_hs300 = top_50_percent
    # print("开始合并")
    df_000300 = get_pro().index_weight(index_code='000016.SH', start_date=cur_000300Date, end_date=cur_000300Date)
    df_000300.rename(columns={'con_code': 'ts_code'}, inplace=True)
    df_000300.rename(columns={'trade_date': 'list_date'}, inplace=True)
    df_000300['symbol'] = df_000300['ts_code'].str.split('.').str[0]
    df_000300.drop(columns=['weight'], inplace=True)

    df_000300 = pd.merge(df_000300, df_sorted[['ts_code', 'i_日均成交', '日均成交额']], on='ts_code', how='left')
    df_000300.to_excel(str_filePre + cur_000300Date + '-000300.xlsx')

    # 获取 dataframe1 中的所有列
    columns_dataframe1 = df_result_hs300.columns

    # 创建缺失列并填充为 '*'
    for column in columns_dataframe1:
        if column not in df_000300.columns:
            df_000300[column] = '**'

    # 将 dataframe2 的列顺序调整为与 dataframe1 一致
    df_000300 = df_000300[columns_dataframe1]

    # 追加 dataframe2 到 dataframe1
    result = pd.concat([df_result_hs300, df_000300], ignore_index=True)

    result.to_excel(str_filePre + '沪深300.xlsx')

    cur_000300_index = get_pro().index_weight(index_code='000300.SH', start_date=cur_000300Date,
                                              end_date=cur_000300Date)
    cur_000300_index.to_excel(str_filePre + '沪深300_本期.xlsx', index=False)


# strStartDate: 周期开启时间
# strEndDate: 周期结束时间
# str_filePre: 输出结果到文件夹
# pre_000300: 上一期沪深300成分股
# cur_000300: 本期沪深300成分股
if __name__ == '__main__':
    # """2024年四季度"""
    strStartDate = '20231101'
    strSeasonDate = '20230731'
    strEndDate = '20241031'
    str_filePre = 'zfile/上证50_2024_4季度/'
    pre_000300 = '20240628'  # 上一期沪深300
    cur_index_date = '20241231'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate,
               # strSeasonDate=strSeasonDate,
               str_filePre=str_filePre,
               cur_000300Date=cur_index_date,
               pre_000300Date=pre_000300)

    # # """2024年4季度"""

    # strStartDate = '20231101'
    # strEndDate = '20241031'
    # str_filePre = 'zfile/沪深300_2024季度4/'
    # pre_000300 = '20240930'  # 上一期沪深300 将上一期的所有股票进入日均总市值样本
    # cur_000300 = '20241231'  # 本期沪深300
    # calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre,
    #            cur_000300Date=cur_000300,
    #            pre_000300Date=pre_000300)


